#include <iostream>
#include <opencv2/core/core.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/highgui/highgui.hpp>
#include "vortex/engine/simple_infer_engine.h"
#include "vortex/utils/timer.h"

namespace vortex
{

constexpr float IMAGENET_DEFAULT_MEAN[3] = {0.485, 0.456, 0.406};
constexpr float IMAGENET_DEFAULT_STD[3] = {0.229, 0.224, 0.225};

    class FasterNet : public SimpleInferEngine
    {
    private:
        BlobInfo m_InputInfo;
        BlobInfo m_OutputInfo;
    public:
        FasterNet(const std::string& engine_path)
        {
            m_InputInfo.name = "input";
            m_InputInfo.shape = {1, 3, 224, 224};

            m_OutputInfo.name = "output";
            m_OutputInfo.shape = {1, 1000, 1, 1};

            this->LoadEngine(engine_path, m_InputInfo, m_OutputInfo);
        }

        virtual void Preprocess(cv::Mat& image) override
        {
            // resize, bgr2rgb, normalize
            cv::Mat temp;
            uint32_t width = m_InputInfo.shape[3];
            uint32_t height = m_InputInfo.shape[2];
            cv::resize(image, temp, cv::Size(width, height));
            uint32_t image_area = width * height;

            float* input_buffer = m_InputBlob->data;
            float* pBlue = input_buffer;
            float* pGreen = input_buffer + image_area;
            float* pRed = input_buffer + image_area * 2;
            unsigned char* pImage = temp.data;
            for (uint32_t i = 0; i < image_area; ++i)
            {
                pRed[i] = (pImage[3 * i + 0] / 255.0 - IMAGENET_DEFAULT_MEAN[0]) / IMAGENET_DEFAULT_STD[0];
                pGreen[i] = (pImage[3 * i + 1] / 255.0 - IMAGENET_DEFAULT_MEAN[1]) / IMAGENET_DEFAULT_STD[1];
                pBlue[i] = (pImage[3 * i + 2] / 255.0 - IMAGENET_DEFAULT_MEAN[2]) / IMAGENET_DEFAULT_STD[2];
            }
        }
    };
}


int main()
{
    std::string engine_path("../fasternet.engine");
    vortex::FasterNet model(engine_path);
    std::string image_path("../../../data/29bb3ece3180_11.jpg");
    cv::Mat image = cv::imread(image_path);

    std::vector<float> outputs;

    vortex::Timer timer;
    
    for (uint32_t count = 0; count <= 50; ++count)
    {
        timer.Start();
        model.Infer(image, outputs);
        float duration = timer.Stop();
        std::cout << "inference time: " << duration << std::endl;
    }

    for (int i = 0; i < 10; ++i)
    {
        std::cout << outputs[i] << std::endl;
    }

    std::cout << "Main\n";
    return 0;
}
